import os
import requests

# 配置路径
SRC_DIR = 'doc-polished'
DST_DIR = 'doc-excercise'
PROMPT_FILE = '练习题提示词.txt'

# 读取提示词内容
with open(PROMPT_FILE, 'r', encoding='utf-8') as f:
    prompt_text = f.read()

# 读取环境变量
DS_BASE_URL = os.getenv('DS_BASE_URL')
DS_API_KEY = os.getenv('DS_API_KEY')

if not DS_BASE_URL or not DS_API_KEY:
    raise RuntimeError("请确保环境变量 DS_BASE_URL 和 DS_API_KEY 已设置。")

headers = {
    'Authorization': f'Bearer {DS_API_KEY}',
    'Content-Type': 'application/json'
}

def polish_text(text, prompt):
    # 这里构造请求payload，具体根据deepseek接口文档调整
    payload = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": prompt},
            {"role": "user", "content": text}
        ],
        "max_tokens": 8000,
        "temperature": 0.7
    }

    url = f"{DS_BASE_URL}/v1/chat/completions"  # 假设接口路径，实际请替换

    response = requests.post(url, json=payload, headers=headers)
    response.raise_for_status()
    data = response.json()
    # 假设返回内容在 data['choices'][0]['message']['content']
    polished = data['choices'][0]['message']['content']
    return polished

def process_files():
    for root, _, files in os.walk(SRC_DIR):
        for file in files:
            if not file.lower().endswith('.md'):
                continue

            src_path = os.path.join(root, file)
            # 计算相对路径，构造目标路径
            rel_path = os.path.relpath(src_path, SRC_DIR)
            dst_path = os.path.join(DST_DIR, os.path.splitext(rel_path)[0]+'-exercise' + '.md')

            if os.path.exists(dst_path):
                print(f"跳过已存在文件: {dst_path}")
                continue

            # 确保目标目录存在
            os.makedirs(os.path.dirname(dst_path), exist_ok=True)

            # 读取原文
            with open(src_path, 'r', encoding='utf-8') as f:
                original_text = f.read()

            print(f"处理文件: {src_path}")
            try:
                polished_text = polish_text(original_text, prompt_text)
            except Exception as e:
                print(f"调用LLM接口失败: {e}")
                continue

            # 写入目标文件
            with open(dst_path, 'w', encoding='utf-8') as f:
                f.write(polished_text)

            print(f"已生成: {dst_path}")

if __name__ == '__main__':
    process_files()
